package com.zelix.analy;

import com.arthas.client.ArthasClient;
import com.arthas.client.ArthasTool;
import me.coley.analysis.SimAnalyzer;
import me.coley.analysis.SimInterpreter;
import me.coley.analysis.value.AbstractValue;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;
import org.objectweb.asm.tree.*;
import org.objectweb.asm.tree.analysis.AnalyzerException;
import org.objectweb.asm.tree.analysis.Frame;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.Enumeration;
import java.util.HashMap;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

/**
 * ZKM 17.0.4 解析并转换成实际指令的例子
 */
public class analyZkm {

    //缓存key,无法动态拦截key的可以设置缓存。动态也可以使用缓存。
    static HashMap<String, String> keyMap = new HashMap() {{
        put("com/zelix/ZKM.main([Ljava/lang/String;)V", "105391775069500");
    }};

    public static void main(String[] args) throws IOException {
        JarFile jarFile = null;
        try {
            jarFile = new JarFile("./file/ZKM.jar");
        } catch (IOException e) {
            e.printStackTrace();
        }

        ArthasClient client = ArthasTool.getClient();

        if (jarFile == null || client == null) {
            return;
        }

        String rule = "com/zelix/ZKM";

        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (entry.getName().equals(rule + ".class")) {//只分析class文件
                //去除文件名后缀 ".class"
                analy(client, jarFile, rule);
            }
        }

        client.disconnect();
    }

    public static void analy(ArthasClient client, JarFile jarFile, String cls) throws IOException {

        System.out.println("***********" + cls + "***********");

        MethodNode enHandle = null;
        MethodNode enMethod = null;
        MethodNode enField = null;

        //获取类所在的classloader,idea中的一些方法不加classloader去运行,可能找不到类
        String classLoader = ArthasTool.getClassLoader(client, cls);
        if (classLoader == null) {
            System.out.println("类 " + cls + " 还未加载,请点击相应功能触发其加载.");
            return;
        }


        /**
         * 两种获取计算key值的方式,一种是通过字符串加密函数,一种是通过获取反射方法的方法,以防解析的方法中只有一种方式
         */
        //字符串解密的ognl表示方式
        String stringMethodFormat = null;
        //监听字符串解密方法
        String watchStringMethodFormat = null;
        //获取反射Method的ognl表示方式
        String reflectMethodFormat = "ognl '@%s@%s(%sL,%sL).toString' -c " + classLoader;
        //获取asm描述Method的ognl表示方式
        String asmMethodFormat = "ognl '#method=@%s@%s(%sL,%sL),#des=@jdk.internal.org.objectweb.asm.Type@getType(#method),#ower=@jdk.internal.org.objectweb.asm.Type@getType(#method.getDeclaringClass()),#ower+\":\"+#method.getName()+\":\"+#des' -c " + classLoader;
        //获取asm描述Field的ognl表示方式
        String asmFieldFormat = "ognl '#field=@%s@%s(%sL,%sL),#des=@jdk.internal.org.objectweb.asm.Type@getType(#field.getType()),#ower=@jdk.internal.org.objectweb.asm.Type@getType(#field.getDeclaringClass()),#ower+\":\"+#field.getName()+\":\"+#des' -c " + classLoader;
        //监听获取反射Method的方法
        String watchReflectMethodFormat = "watch %s %s {params[1]} 'params.length==2 && params[0] instanceof java.lang.Long && params[1] instanceof java.lang.Long && returnObj !=null && returnObj instanceof java.lang.reflect.Method && params[0]==%sL' -n 1";
        //监听获取反射Field的方法
        String watchReflectFieldFormat = "watch %s %s {params[1]} 'params.length==2 && params[0] instanceof java.lang.Long && params[1] instanceof java.lang.Long && returnObj !=null && returnObj instanceof java.lang.reflect.Field && params[0]==%sL' -n 1";


        ZipEntry entry = jarFile.getEntry(cls + ".class");

        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, ClassReader.EXPAND_FRAMES);

            List<MethodNode> methods = cn.methods;

            //查找符合本类的字符串解密方法并初始化解密方法,字符串加密都在本类
            for (MethodNode methodNode : methods) {
                if (methodNode.desc.equals("(IJ)Ljava/lang/String;")) {
                    stringMethodFormat = "ognl '@" + cls.replaceAll("/", ".") + "@" + methodNode.name + "(%d,%sL^%sL)' -c " + classLoader;
                    watchStringMethodFormat = "watch " + cls.replaceAll("/", ".") + " " + methodNode.name + " {params[1]} 'params.length==2 && params[0] instanceof java.lang.Integer && params[1] instanceof java.lang.Long && returnObj !=null && returnObj instanceof String && params[0]==%d' -n 1";
                }
            }

            //遍历方法
            for (MethodNode methodNode : methods) {
                System.out.println("------------" + methodNode.name + methodNode.desc + "------------");
                //计算long值,遍历每个方法前初始化,不同方法值不同
                String key = keyMap.get(cls + "." + methodNode.name + methodNode.desc);

                //frame分析
                SimAnalyzer analyzer = new SimAnalyzer(new SimInterpreter());
                analyzer.setThrowUnresolvedAnalyzerErrors(false);
                Frame<AbstractValue>[] frames = null;
                try {
                    frames = analyzer.analyze(cn.name, methodNode);
                } catch (AnalyzerException e) {
                    e.printStackTrace();
                    break;
                }


                InsnList srcList = methodNode.instructions;

                //遍历指令
                for (int i = 0; i < srcList.size(); i++) {

                    //查找invokedynamic指令
                    if (srcList.get(i) instanceof InvokeDynamicInsnNode) {
                        InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) srcList.get(i);
                        if (dynamicNode.desc.equals("(IJ)Ljava/lang/String;") && dynamicNode.bsm.getOwner().equals(cls)) {
                            //字符串解密,函数描述为(IJ)Ljava/lang/String;
                            /**
                             * SIPUSH 14142
                             * LDC 792218991624472234L
                             * LLOAD 1
                             * LXOR
                             * INVOKEDYNAMIC j (IJ)Ljava/lang/String; handle[H_INVOKESTATIC com/zelix/ZKM.b(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                             */
                            //System.out.println("字符串加密:" + dynamicNode.name + ":" + dynamicNode.desc);

                            Frame<AbstractValue> frame = frames[i];
                            int size = frame.getStackSize();
                            //参数栈必须大于2个,主要看倒数后两个,倒数第一个是I,倒数第二个是J
                            if (size >= 2) {
                                //倒数第一个是I,倒数第二个是J
                                if (frame.getStack(size - 2).getType() == Type.INT_TYPE && frame.getStack(size - 1).getType() == Type.LONG_TYPE) {

                                    //获取倒数第二个参数J计算而来的指令
                                    List<AbstractInsnNode> insns = frame.getStack(size - 1).getInsns();
                                    //命令必须是三个以上参数
                                    if (insns.size() >= 3) {
                                        //最后一个命令需要是LXOR,最后需要LXOR进行异或运算的
                                        if (insns.get(insns.size() - 1) instanceof InsnNode && insns.get(insns.size() - 1).getOpcode() == Opcodes.LXOR) {
                                            //中间命令为类似LLOAD 4,这个变量可能是多个变量计算而来,所以会展开很多命令,所以只判断一个
                                            if (insns.get(insns.size() - 2) instanceof VarInsnNode && insns.get(insns.size() - 2).getOpcode() == Opcodes.LLOAD) {
                                                //再判断第一个是不是LDC命令
                                                if (insns.get(0) instanceof LdcInsnNode) {

                                                    LdcInsnNode ldcInsnNode = (LdcInsnNode) insns.get(0);
                                                    String longValue = ldcInsnNode.cst.toString();

                                                    String intValue = null;
                                                    AbstractInsnNode intInsnNode = null;
                                                    //倒数第二个参数符合判断的情况下获取倒数第一个参数
                                                    //已经判断过倒数第一个参数必须是I
                                                    if (frame.getStack(size - 2).getValue() != null) {
                                                        intValue = frame.getStack(size - 2).getValue().toString();
                                                        //一般只有一个指令
                                                        intInsnNode = frame.getStack(size - 2).getInsns().get(0);
                                                    }

                                                    if (intValue != null) {//值解析成功
                                                        if (key != null) {
                                                            String rtn = "";
                                                            rtn = client.sendCommand(String.format(stringMethodFormat, Integer.valueOf(intValue), longValue, key));
                                                            //System.out.println("字符串解密:" + ArthasTool.removeStringLebal(rtn));

                                                            //获取指令索引
                                                            //SIPUSH
                                                            //System.out.println(srcList.indexOf(intInsnNode));
                                                            //LDC
                                                            //System.out.println(srcList.indexOf(insns.get(0)));
                                                            //LLOAD
                                                            //System.out.println(srcList.indexOf(insns.get(insns.size() - 2)));
                                                            //LXOR
                                                            //System.out.println(srcList.indexOf(insns.get(insns.size() - 1)));
                                                            //当前InvokeDynamicInsnNode索引
                                                            //System.out.println(i);

                                                            //替换LDC指令
                                                            srcList.set(srcList.get(i), new LdcInsnNode(ArthasTool.removeStringLebal(rtn)));
                                                            //其余替换NOP指令,保持指令列表平衡
                                                            srcList.set(insns.get(insns.size() - 1), new InsnNode(Opcodes.NOP));
                                                            srcList.set(insns.get(insns.size() - 2), new InsnNode(Opcodes.NOP));
                                                            srcList.set(insns.get(0), new InsnNode(Opcodes.NOP));
                                                            srcList.set(intInsnNode, new InsnNode(Opcodes.NOP));
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }

                        } else if (dynamicNode.desc.contains("JJ)")) {
                            //方法属性解密,函数最少两个参数,这两个必须是最后两个,而且必须是long
                            /**
                             * LDC -5040441400421082001L
                             * LLOAD 1
                             * INVOKEDYNAMIC k (Ljava/lang/Object;JJJ)Ljava/lang/StringBuilder; handle[H_INVOKESTATIC com/zelix/v99.a(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                             */
                            //System.out.println("方法加密:" + dynamicNode.name + ":" + dynamicNode.desc);

                            if (enHandle == null) {
                                //加密逻辑处理方法需要返回一个MethodHandle对象
                                enHandle = getMethod(jarFile, dynamicNode.bsm.getOwner(), ")Ljava/lang/invoke/MethodHandle;");
                                //加密方法
                                enMethod = getMethod(jarFile, dynamicNode.bsm.getOwner(), "(JJ)Ljava/lang/reflect/Method;");
                                //加密属性
                                enField = getMethod(jarFile, dynamicNode.bsm.getOwner(), "(JJ)Ljava/lang/reflect/Field;");
                            }
                            if (enHandle != null) {
                                //System.out.println("加密逻辑处理方法为:" + dynamicNode.bsm.getOwner() + ":" + enHandle.name + ":" + enHandle.desc);
                                //System.out.println("加密方法为:" + dynamicNode.bsm.getOwner() + ":" + enMethod.name + ":" + enMethod.desc);
                                //System.out.println("加密属性方法为:" + dynamicNode.bsm.getOwner() + ":" + enField.name + ":" + enField.desc);
                            }

                            Frame<AbstractValue> frame = frames[i];
                            int size = frame.getStackSize();
                            //参数栈必须大于2个
                            if (size >= 2) {
                                //第一个参数必须是J,第二个参数必须是J
                                if (frame.getStack(size - 2).getType() == Type.LONG_TYPE && frame.getStack(size - 1).getType() == Type.LONG_TYPE) {
                                    //获取第二个参数J计算而来的指令
                                    List<AbstractInsnNode> insns = frame.getStack(size - 1).getInsns();
                                    //判断最后一个指令是不是LLOAD
                                    if (insns.get(insns.size() - 1) instanceof VarInsnNode && insns.get(insns.size() - 1).getOpcode() == Opcodes.LLOAD) {
                                        //第二个参数符合判断的情况下获取第一个参数
                                        String longValue = null;
                                        //已经判断过第一个参数必须是J,这里是快速获取值,可以通过 insns = frame.getStack(size - 2).getInsns();获取指令,然后判断LdcInsnNode指令去获取,这里因为模式固定,所以直接获取值了
                                        if (frame.getStack(size - 2).getValue() != null) {
                                            longValue = frame.getStack(size - 2).getValue().toString();
                                            AbstractInsnNode longInsnNode = frame.getStack(size - 2).getInsns().get(0);


                                            //获取指令索引
                                            //LDC
                                            //System.out.println(srcList.indexOf(longInsnNode));
                                            //LLOAD
                                            //System.out.println(srcList.indexOf(insns.get(insns.size() - 1)));
                                            //当前InvokeDynamicInsnNode索引
                                            //System.out.println(i);


                                            //替换对应指令
                                            //区分方法和属性的值
                                            int tag = getTagInt(dynamicNode.name, Long.valueOf(key));
                                            if (tag == 110 || tag == 112 || tag == 115 || tag == 109) {
                                                //属性
                                                String rtn = "";
                                                //rtn = client.sendCommand(String.format(reflectMethodFormat, dynamicNode.bsm.getOwner().replaceAll("/", "."), enField.name, longValue, key));
                                                //System.out.println("属性加密:" + rtn);
                                                //获取指令对象描述
                                                rtn = client.sendCommand(String.format(asmFieldFormat, dynamicNode.bsm.getOwner().replaceAll("/", "."), enField.name, longValue, key));
                                                rtn = asmdes(rtn);
                                                //System.out.println("ASM对象描述" + rtn);
                                                String[] des = rtn.split(":");

                                                int opcodes = -1;
                                                if (tag == 110) {
                                                    opcodes = Opcodes.GETFIELD;
                                                } else if (tag == 112) {
                                                    opcodes = Opcodes.PUTFIELD;
                                                } else if (tag == 115) {
                                                    opcodes = Opcodes.GETSTATIC;
                                                } else {
                                                    opcodes = Opcodes.PUTSTATIC;
                                                }
                                                srcList.set(srcList.get(i), new FieldInsnNode(opcodes, des[0], des[1], des[2]));


                                            } else {
                                                //方法
                                                String rtn = "";
                                                //rtn = client.sendCommand(String.format(reflectMethodFormat, dynamicNode.bsm.getOwner().replaceAll("/", "."), enMethod.name, longValue, key));
                                                //System.out.println("方法加密:" + rtn);
                                                //获取指令对象描述
                                                rtn = client.sendCommand(String.format(asmMethodFormat, dynamicNode.bsm.getOwner().replaceAll("/", "."), enMethod.name, longValue, key));
                                                rtn = asmdes(rtn);
                                                //System.out.println("ASM对象描述" + rtn);
                                                String[] des = rtn.split(":");

                                                int opcodes = -1;
                                                if (tag == 111) {
                                                    opcodes = Opcodes.INVOKEVIRTUAL;
                                                } else if (tag == 108) {
                                                    opcodes = Opcodes.INVOKESTATIC;
                                                } else {
                                                    opcodes = Opcodes.INVOKESPECIAL;
                                                }
                                                srcList.set(srcList.get(i), new MethodInsnNode(opcodes, des[0], des[1], des[2]));

                                            }
                                            //其余替换NOP指令
                                            srcList.set(insns.get(insns.size() - 1), new InsnNode(Opcodes.NOP));
                                            srcList.set(longInsnNode, new InsnNode(Opcodes.NOP));

                                        }
                                    }
                                }
                            }
                        } else if (dynamicNode.bsm.getOwner().equals("java/lang/invoke/LambdaMetafactory")) {
                            System.out.println("Lambda表达式:" + dynamicNode.name + ":" + dynamicNode.desc);
                        } else {
                            System.out.println("未知用途:");
                            System.out.println("实际方法描述:" + dynamicNode.name + ":" + dynamicNode.desc);
                            System.out.println("调用方法描述:" + dynamicNode.bsm.getOwner() + ":" + dynamicNode.bsm.getName());
                        }
                    }
                }

                //去掉填充的NOP指令
                //再次遍历指令,正序后续nop不会处理,先处理最后的
                for (int i = srcList.size() - 1; i > 0; i--) {
                    if (srcList.get(i) instanceof InsnNode && srcList.get(i).getOpcode() == Opcodes.NOP) {
                        srcList.remove(srcList.get(i));
                    }
                }


                methodNode.visitEnd();
            }

            //写入文件
            ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
            cn.accept(cw);
            Files.write(Paths.get(cls.replaceAll("/", ".") + ".class"), cw.toByteArray());
        }
    }

    /**
     * 根据方法描述获取对应方法
     *
     * @param jarFile
     * @param cls
     * @return
     * @throws IOException
     */
    public static MethodNode getMethod(JarFile jarFile, String cls, String ends) throws IOException {
        ZipEntry entry = jarFile.getEntry(cls + ".class");

        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, ClassReader.EXPAND_FRAMES);

            List<MethodNode> methods = cn.methods;

            //遍历方法
            for (MethodNode methodNode : methods) {
                if (methodNode.desc.endsWith(ends)) {
                    return methodNode;
                }
            }
        }
        return null;
    }

    /**
     * INVOKEDYNAMIC命令的名字转换为区分属性和方法的int值,解析返回值为MethodHandle的方法得到的
     *
     * @param tag
     * @param l
     * @return
     */
    public static int getTagInt(String tag, long l) {
        int i = (int) l;
        i = i & 7;
        i = tag.charAt(0) ^ i;
        return i;
    }

    /**
     * 返回ower+name+des的形式
     *
     * @param rtn
     * @return
     */
    public static String asmdes(String rtn) {
        rtn = ArthasTool.removeStringLebal(rtn);
        String[] split = rtn.split(":");
        if (split.length == 3) {
            if (split[0].startsWith("L") && split[0].endsWith(";")) {//owner应该是对象格式
                return split[0].substring(1, split[0].length() - 1) + ":" + split[1] + ":" + split[2];
            }
        }
        return null;
    }
}

